#!/usr/bin/env python

# Copyright 2012 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""OpenID Connect Client Demo Application."""

import os
import random
import urllib
from google.appengine.ext import webapp
from google.appengine.ext import db
from google.appengine.ext.webapp.util import run_wsgi_app
from google.appengine.api import memcache
from google.appengine.ext.webapp import template
import logging
from oauth2client.client import OAuth2Credentials, flow_from_clientsecrets
from oauth2client.client import FlowExchangeError
from openidconnect.client import openidconnect_flow_from_clientsecrets
from openidconnect.error import FlowUserInfoError, FlowTokenInfoError
from streetidentity_rp.client import StreetIdentityFlow, \
    StreetIdentityCredentials, AttributeCredentials, \
    StreetIdentityAttributeProviderFlow, ExtendedStreetIdentityCredentials, \
    AttributeType
from streetidentity_rp.error import FlowStreetIdentityError
from streetidentity_rp.oauth2client_ext import ClientCredentials

__author__ = "Maciej Machulak"
__maintainer__ = "Maciej Machulak"
__email__ = "mmachulak@google.com"

__copyright__ = "Copyright 2012 Google Inc. All Rights Reserved."
__license__ = "Apache License 2.0"
__version__ = "0.2"
__status__ = "Prototype"

SESSION_ID = 'session_id'
CALLBACK_URL = '/oauth2callback'
STREETIDENTITY_ENTRY_URL = '/streetidentity'
HEADER_SET_COOKIE = 'Set-Cookie'

OPENIDCONNECT_FLOW = openidconnect_flow_from_clientsecrets(
    'config/client_secrets.json')

IDP_FLOW = StreetIdentityFlow(
    client_id=OPENIDCONNECT_FLOW.client_id,
    client_secret=OPENIDCONNECT_FLOW.client_secret,
    attribute_address_uri=
    'https://streetidentity-console.appspot.com/api/v1/attribute_info/',
    attribute_age_uri=
    'https://streetidentity-console.appspot.com/api/v1/attribute_info/age',
    attribute_gender_uri=
    'https://streetidentity-console.appspot.com/api/v1/attribute_info/gender',
    attribute_token_uri=
    'https://streetidentity-console.appspot.com/api/v1/attribute_token')

PROVIDER_DEMO_AP = "776695651878.apps.googleusercontent.com"

KNOWN_ATTRIBUTE_PROVIDERS = [PROVIDER_DEMO_AP]

AP_CREDENTIALS = flow_from_clientsecrets(
    filename='config/ap_client_secrets.json',
    scope='ap_defined_scope')

PROVIDER_CREDENTIALS = {
    PROVIDER_DEMO_AP: ClientCredentials(
        client_id=AP_CREDENTIALS.client_id,
        client_secret=AP_CREDENTIALS.client_secret,
        token_uri=AP_CREDENTIALS.token_uri)}

PROVIDER_FLOWS = {
    PROVIDER_DEMO_AP: StreetIdentityAttributeProviderFlow(
        "https://streetidentity-ap-client.appspot.com/api/streetidentity/")
}


def login_required(handler):
    """Requires that a user is logged in."""

    def check_login(self, *args):
        try:
            session_id = self.request.cookies[SESSION_ID]
            logging.debug('session id: %s' % session_id)
        except KeyError:
            session_id = None
        if session_id is None:
            logging.debug('session id is none - redirecting to main page')
            self.redirect('/')
            return
        session = memcache.get(session_id)
        if session is None:
            logging.error(
                'session id found but session could not be retrieved from the'
                ' database')
            self.response.headers.add_header(HEADER_SET_COOKIE,
                SESSION_ID + '=')
            logging.error('deleted session cookie')
            self.redirect('/')
            return
        self.credentials = session
        logging.debug('user has been logged in successfully')
        handler(self, *args)

    return check_login


class AuthNError(Exception):
    """ Raised when there is an error with the login component. """

    def __init__(self, msg):
        Exception.__init__(self)
        self.msg = msg

    def __str__(self):
        return repr(self.msg)


class User(db.Model):
    """ Represents the user of the application. """

    id = db.StringProperty(required=True)
    user_info = db.Blob()
    address = db.Blob()
    phone = db.Blob()


class Session(db.Model):
    """ Represents the active session of the application. """

    id = db.StringProperty(required=True)
    user = db.ReferenceProperty(User, required=True)


class MainHandler(webapp.RequestHandler):
    """ Handles the main entry to the application. """

    def get(self, *args):
        path = os.path.join(os.path.dirname(__file__), 'templates/index.html')
        self.response.out.write(template.render(path, {}))


class SignInHandler(webapp.RequestHandler):
    """ Generates a sign in link for the correct identity provider. """

    def get(self, *args):
        callback = self.request.relative_url(CALLBACK_URL)
        logging.debug('Callback: %s' % callback)
        authorize_url = IDP_FLOW.step1_get_authorize_url(callback)
        self.redirect(authorize_url)


class SignOutHandler(webapp.RequestHandler):
    """ Destroys the session for the user. """

    def get(self, *args):
        session_id = self.request.cookies.get(SESSION_ID, '')
        if session_id:
            self.response.headers.add_header(HEADER_SET_COOKIE,
                SESSION_ID + '=')
            memcache.delete(session_id)
        path = os.path.join(os.path.dirname(__file__), 'templates/index.html')
        self.response.out.write(template.render(path, {}))


class AccountHandler(webapp.RequestHandler):
    """ Account handler for the user - user must be authenticated. """

    @login_required
    def get(self, *args):
        if len(self.credentials.attribute_list):
            values = {
                'userinfo': self.credentials.userinfo,
                'userinfo_string': repr(self.credentials.userinfo),
                'attributes': self.credentials.attribute_list,
                'attributes_string': repr(self.credentials.attribute_list)
            }
        else:
            values = {
                'userinfo': self.credentials.userinfo,
                'userinfo_string': repr(self.credentials.userinfo)
            }

        path = os.path.join(os.path.dirname(__file__),
            'templates/account.html')
        self.response.out.write(template.render(path, values))


class OAuthHandler(webapp.RequestHandler):
    """ OAuth Redirect handler """

    def get(self, *args):
        logging.debug('exchanging code for access token')
        try:
            # STEP 2: exchange authorization for access token
            logging.debug(
                '---> STEP 2: EXCHANGE AUTHORIZATION FOR ACCESS TOKEN')
            credentials = IDP_FLOW.step2_exchange(self.request.params)
        except FlowExchangeError:
            logging.error(
                'error during OAuth dance - could not exchange grant for the '
                'token')
            values = {
                'error_message': 'could not exchange grant for the token'
            }
            path = os.path.join(os.path.dirname(__file__),
                'templates/error.html')
            self.response.out.write(template.render(path, values))
            return

        logging.debug('verifying received token')
        try:
            # STEP 3: verify access token
            logging.debug(
                '---> STEP 3: VERIFY ACCESS TOKEN')
            verified_token_credentials = IDP_FLOW.step3_verify_access_token(
                credentials)
        except FlowTokenInfoError:
            logging.error('received token is not valid')
            values = {
                'error_message': 'could not validate the token received from '
                                 'the IDP'
            }
            path = os.path.join(os.path.dirname(__file__),
                'templates/error.html')
            self.response.out.write(template.render(path, values))
            return

        logging.debug(
            'using access token to access user information from the IDP')
        try:
            # STEP 3: get user info
            logging.debug('---> STEP 4: GETTING USER INFO')
            openidconnect_credentials = IDP_FLOW.step4_userinfo(
                verified_token_credentials)
            logging.debug('obtained user info')
            key = str(random.random())[2:] + str(random.random())[2:]
            memcache.set(key, openidconnect_credentials)
            logging.debug('user info stored in memcache successfully')

            # STEP 5: get a list of all attributes
            logging.debug('---> STEP 5: GETTING ATTRIBUTE INFORMATION')
            attribute_list_info = IDP_FLOW.step5_get_attribute_information(
                openidconnect_credentials,
                attribute_type_list=
                [AttributeType.STREET_ADDRESS,
                 AttributeType.AGE,
                 AttributeType.GENDER])
            logging.debug(
                'obtained a list of attributes: %s' % len(attribute_list_info))

            # STEP 6: get tokens for all discovered attributes
            logging.debug(
                '---> STEP 6: GETTING TOKENS FOR DISCOVERED ATTRIBUTES')
            attribute_credentials_list = []

            for attribute_info in attribute_list_info.attribute_list:
                logging.debug('checking attribute info %s' % attribute_info)
                # check if the attribute provider is known (i.e. there is a
                # relationship with this provider
                if attribute_info.provider_id in KNOWN_ATTRIBUTE_PROVIDERS:
                    logging.debug(
                        'trying to get a token for known attribute provider: ' \
                        '%s' % attribute_info.provider_id)
                    authorized_attribute_info = IDP_FLOW.step6_get_token(
                        openidconnect_credentials, attribute_info)
                    attribute_credentials = AttributeCredentials(
                        PROVIDER_CREDENTIALS[attribute_info.provider_id],
                        authorized_attribute_info)
                    attribute_credentials_list.append(attribute_credentials)
                else:
                    logging.debug(
                        'provider unknown: %s' % attribute_info.provider_id)

            logging.debug('got authorization to access %s attributes' % len(
                attribute_credentials_list))
            streetidentity_credentials = StreetIdentityCredentials(
                openidconnect_credentials,
                attribute_credentials_list)
            memcache.set(key, streetidentity_credentials)
            logging.debug(
                'street identity credentials stored in memcache successfully')

            # STEP 7: Getting attributes from providers
            logging.debug("---> STEP 7: GETTING ATTRIBUTES FROM APs")
            attributes = []

            for attribute_credentials in attribute_credentials_list:
                logging.debug(attribute_credentials)
                oauth_credentials = attribute_credentials.oauth_credentials
                logging.debug(oauth_credentials)
                authorized_attribute_info = attribute_credentials\
                .authorized_attribute_info
                ap_flow = PROVIDER_FLOWS[
                          authorized_attribute_info.attribute_info.provider_id]
                attribute = ap_flow.step7_get_attribute(oauth_credentials,
                    authorized_attribute_info.attribute_info.id,
                    authorized_attribute_info.access_token.access_token)
                logging.debug(
                    'successfully obtained attribute from AP: %s' % attribute)
                attributes.append(attribute)

            logging.debug(
                'obtained a list of attributes from AP: %s' % len(attributes))
            extended_streetidentity_credentials = \
            ExtendedStreetIdentityCredentials(
                streetidentity_credentials,
                attributes)
            memcache.set(key, extended_streetidentity_credentials)
            logging.debug(
                'extended street identity credentials stored in memcache '
                'successfully')

            self.response.headers.add_header(HEADER_SET_COOKIE,
                SESSION_ID + '=' + key)
            self.redirect('/account')
        except FlowUserInfoError:
            logging.error(
                'error during OpenID Connect flow - could not retrieve user '
                'info')
            values = {
                'error_message': 'user info could not be retrieved'
            }
            path = os.path.join(os.path.dirname(__file__),
                'templates/error.html')
            self.response.out.write(template.render(path, values))
            return
        except FlowStreetIdentityError:
            logging.error(
                'error during Street Identity flow - could not retrieve the list of attributes')
            values = {
                'error_message': 'attribute list could not be retrieved'
            }
            path = os.path.join(os.path.dirname(__file__),
                'templates/error.html')
            self.response.out.write(template.render(path, values))

routes = [('/', MainHandler),
    ('/sign_in', SignInHandler),
    ('/sign_out', SignOutHandler),
    ('/account', AccountHandler),
    (CALLBACK_URL, OAuthHandler)]

application = webapp.WSGIApplication(routes, debug=True)

def main():
    logging.getLogger().setLevel(logging.DEBUG)
    run_wsgi_app(application)

if __name__ == "__main__":
    main()